#include <winsock2.h>
#include <ws2tcpip.h>
#include <Windows.h>

#include <TlHelp32.h>

#include "../Common/def.h"
#include "../Common/xConfig.h"
#include "../Common/xFile.h"
#include "../Common/Hook.h"

#include "HookProxy.h"

#include <stdlib.h>

#pragma comment(lib, "ws2_32.lib")

Config g_cfg;
std::map<DWORD, BrigeCode> g_bridge_code;
HHOOK g_hhk;

typedef int (WSAAPI *connect_call)(SOCKET ,const struct sockaddr*, int);
connect_call connect_true = NULL;

bool do_connect_hook();
int WSAAPI connect_hook(SOCKET s, const struct sockaddr FAR *name, int namelen);
bool undo_connect_hook(DWORD addr);

bool do_connect_hook(){
    HINSTANCE conn_inst = GetModuleHandle(WINSOCK2);
	if(!conn_inst){
		return false;
	}

	connect_true = (connect_call)GetProcAddress(conn_inst, "connect");
	PBYTE obj_func = (PBYTE)connect_true;
	PBYTE hook_func = (PBYTE)connect_hook;
	const int code_len = 5;
	obj_func = skip_jmp(obj_func);
	hook_func = skip_jmp(hook_func);

	DWORD cur_pid = GetCurrentProcessId();
	PBYTE try_addr = alloc_virtual_addr(obj_func);
	if(!try_addr) return false;


	BYTE orig_code[64], new_code[64];
	PBYTE pcode = orig_code;
	memcpy(pcode, obj_func, code_len);
	pcode += code_len;

	LONG dist = obj_func - try_addr;
	if(dist < 0) dist = -dist;

	// dist < 2G
	if(dist >= 0x40000000){
		return false;
	} else {
		pcode[0] = 0xe9; // jmp
		pcode ++;
		*(DWORD *)pcode = (DWORD)obj_func  - (DWORD)try_addr - 5;
		memcpy(try_addr, orig_code, code_len + 5);
	}	

	for(int i=0; i<code_len; i++) new_code[i] = 0x90; // nop
	new_code[0] = 0xe9;
	pcode = new_code + 1;

	if(dist >= 0x40000000){
		return false;
	} else {
		*(LONG *)pcode = (LONG)hook_func - (LONG)obj_func - 5;
		WriteProcessMemory(INVALID_HANDLE_VALUE, obj_func, new_code, code_len, 0);
	}
	BrigeCode bg;
	bg.orig_addr = obj_func;
	bg.new_addr = try_addr;
	string bridge_code = (char *)orig_code;
	bg.code = bridge_code.substr(0, code_len);
	g_bridge_code[(DWORD)try_addr] = bg;

	connect_true = (connect_call)try_addr;
	return true;
}

int WSAAPI connect_hook(SOCKET s, const struct sockaddr FAR *name, int namelen){
	string proxy_addr = g_cfg.proxy_addr;
	string proxy_ip = "127.0.0.1";
	unsigned short proxy_port = 15366;
	int k = proxy_addr.find(":");
	//proxy_ip = proxy_addr.substr(0, k);
	//if(k >= 0) proxy_port = atoi(proxy_addr.substr(k+1).c_str());
	//MessageBox(NULL, proxy_ip.c_str(), 0, MB_OK);

	sockaddr_in proxy_name;
	proxy_name.sin_family = AF_INET;
	proxy_name.sin_addr.s_addr = inet_addr(proxy_ip.c_str());
	proxy_name.sin_port = htons(proxy_port);
	return connect_true(s, (const struct sockaddr *)&proxy_name, sizeof(sockaddr_in));
}

bool undo_connect_hook(DWORD addr){
	if(g_bridge_code.count(addr) == 0) return false;

	const BrigeCode& bg = g_bridge_code[addr];
	PBYTE orig_addr = bg.orig_addr;
	PBYTE new_addr = bg.new_addr;
	string code = bg.code;

	if(WriteProcessMemory(INVALID_HANDLE_VALUE, orig_addr, code.c_str(), code.size(), 0)){
		return false;
	}
	SYSTEM_INFO sSysInfo =  {0};
	GetSystemInfo(&sSysInfo);
	const DWORD alloc_size = sSysInfo.dwAllocationGranularity;
	if(!VirtualFree(orig_addr, alloc_size,MEM_RELEASE)){
		return false;
	}
	return true;
}

LRESULT WINAPI CallWndProc(int nCode, WPARAM wParam, LPARAM lParam) {
	return CallNextHookEx(g_hhk, nCode, wParam, lParam); 
}

#ifdef __cplusplus
extern "C"{
#endif
EXPORT_API bool hook_start(){
	HINSTANCE hinst = LoadLibrary(TEXT(HOOK_PATH));
	if(!hinst) {
		printf("LoadLibrary %s error\n", HOOK_PATH);
		return false;
	}

	g_hhk = SetWindowsHookEx(WH_CALLWNDPROC, CallWndProc, hinst, 0);
	if(!g_hhk) {
		printf("SetWindowsHookEx error\n");
		return false;
	}
	return true;
}

EXPORT_API bool hook_stop(){
	UnhookWindowsHookEx(g_hhk);
	return true;
}

#ifdef __cplusplus
}
#endif

BOOL WINAPI DllMain(HINSTANCE hinstDLL, DWORD fdwReason, LPVOID lpvReserved){
	char buf[64];
	sprintf(buf, "C:\\process_%d.txt", GetCurrentProcessId());
	ofstream flog(buf);
	flog << "DllMain" << endl;
	if(fdwReason == DLL_PROCESS_ATTACH){
		config_init(CONFIG_PATH, g_cfg);
		do_connect_hook();
	}else if(fdwReason == DLL_PROCESS_DETACH){
		undo_connect_hook((DWORD)connect_true);
	}
	return true;
}
